from itertools import cycle
import re
import sys

def makeCodonTable(codonPath):
    '''Input: Copy/Paste of Wikipedia page on DNA Codons
        ['Amino acid', <-- e.g. Ala/A
        'Codons',
        'Compressed',
        'Amino acid', <-- Second set of columns
        'Codons',
        'Compressed']
        Output: Dictionary of Amino Acid to Cycle Iterator of Codons'''
    codonFile = open(codonPath,'r')
    codonText = codonFile.readlines()
    cycleDict = {}
    for i, line in enumerate(codonText):
        if i>0:
            items = line.split('\t')
            for j in [0,3]:
                cycleDict[items[j].split('/')[-1]] = cycle(re.split(',\s?',items[j+1].strip()))
    codonFile.close()
    return cycleDict

def main(argv):
    '''Input: _protein sequence as string_ [path to codon table file]
        Output: DNA sequence of cyclic codon usage of input protein sequence'''
    if len(argv)<3:
        if len(argv)<2:
            print '''
                Input: _protein sequence as string_ [path to codon table file]
                Output: DNA sequence of cyclic codon usage of input protein sequence
                '''
            return 1
        script, proteinSeq = argv
        codonPath='codons_wiki.tsv'
    else:
        script, proteinSeq, codonPath = argv
    dnaSeq = ''
    proteinSeq = proteinSeq.upper()
    cycleDict = makeCodonTable(codonPath)
    for char in proteinSeq:
        dnaSeq = dnaSeq + cycleDict[char].next()
    return dnaSeq

if __name__ == "__main__":
    sys.exit(main(sys.argv))